package main

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"time"

	"github.com/gorilla/websocket"
)

// 定义消息结构
type Message struct {
	Username  string `json:"username"`
	Content   string `json:"content"`
	Type      string `json:"type"`
	Timestamp string `json:"timestamp"`
}

// 客户端结构
type Client struct {
	conn     *websocket.Conn
	username string
	send     chan Message
}

// 聊天服务器结构
type ChatServer struct {
	clients    map[*Client]bool
	register   chan *Client
	unregister chan *Client
	broadcast  chan Message
}

// 创建新的聊天服务器
func NewChatServer() *ChatServer {
	return &ChatServer{
		clients:    make(map[*Client]bool),
		register:   make(chan *Client, 100),   // 使用带缓冲的通道
		unregister: make(chan *Client, 100),   // 使用带缓冲的通道
		broadcast:  make(chan Message, 10000), // 使用带缓冲的通道
	}
}

// 启动聊天服务器
func (cs *ChatServer) Start() {
	for {
		select {
		case client := <-cs.register:
			cs.clients[client] = true
			log.Printf("客户端 %s 已连接\n", client.username)
			//广播用户加入消息
			cs.broadcast <- Message{
				Username:  "系统",
				Content:   fmt.Sprintf("%s 加入了聊天", client.username),
				Type:      "system",
				Timestamp: time.Now().Format("15:04:05"),
			}
		case client := <-cs.unregister:
			if _, ok := cs.clients[client]; ok {
				delete(cs.clients, client)
				close(client.send)
				log.Printf("客户端 %s 已断开连接\n", client.username)
				// 广播用户离开消息
				cs.broadcast <- Message{
					Username:  "系统",
					Content:   fmt.Sprintf("%s 离开了聊天", client.username),
					Type:      "system",
					Timestamp: time.Now().Format("15:04:05"),
				}
			}
		case message := <-cs.broadcast:
			// 向所有客户端广播消息
			for client := range cs.clients {
				if message.Type == "user" && message.Username == client.username {
					continue
				}
				select {
				case client.send <- message:
				default:
					close(client.send)
					delete(cs.clients, client)
				}
			}
		}
	}
}

// WebSocket 升级器
var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return true // 允许所有来源
	},
}

// 处理WebSocket连接
func (cs *ChatServer) handleWebSocket(w http.ResponseWriter, r *http.Request) {
	// 从URL参数获取用户名
	username := r.URL.Query().Get("username")
	if username == "" {
		http.Error(w, "缺少用户名参数", http.StatusBadRequest)
		return
	}

	// 升级HTTP连接为WebSocket连接
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println("WebSocket升级失败:", err)
		return
	}

	client := &Client{
		conn:     conn,
		username: username,
		send:     make(chan Message, 256),
	}

	// 注册客户端
	cs.register <- client

	// 启动读取和写入goroutine
	go client.writePump()
	go client.readPump(cs)
}

// 读取客户端消息
func (c *Client) readPump(cs *ChatServer) {
	defer func() {
		cs.unregister <- c
		c.conn.Close()
	}()

	c.conn.SetReadLimit(512)
	c.conn.SetReadDeadline(time.Now().Add(60 * time.Second))
	c.conn.SetPongHandler(func(string) error { c.conn.SetReadDeadline(time.Now().Add(60 * time.Second)); return nil })

	for {
		_, message, err := c.conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				log.Printf("读取错误: %v", err)
			}
			break
		}

		// 解析接收到的消息
		var msg Message
		if err := json.Unmarshal(message, &msg); err != nil {
			log.Printf("解析消息失败: %v", err)
			continue
		}

		// 设置消息发送者用户名
		msg.Username = c.username

		// 发送到广播通道
		cs.broadcast <- msg
	}
}

// 向客户端写入消息
func (c *Client) writePump() {
	ticker := time.NewTicker(55 * time.Second)
	defer func() {
		ticker.Stop()
		c.conn.Close()
	}()

	for {
		select {
		case message, ok := <-c.send:
			c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
			if !ok {
				// 通道关闭
				c.conn.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}

			w, err := c.conn.NextWriter(websocket.TextMessage)
			if err != nil {
				return
			}

			// 序列化消息
			messageJSON, err := json.Marshal(message)
			if err != nil {
				log.Printf("序列化消息失败: %v", err)
				w.Close()
				return
			}

			w.Write(messageJSON)

			// 添加上次发送的消息数量
			n := len(c.send)
			for i := 0; i < n; i++ {
				w.Write([]byte("\n"))
				messageJSON, err := json.Marshal(<-c.send)
				if err != nil {
					log.Printf("序列化消息失败: %v", err)
					w.Close()
					return
				}
				w.Write(messageJSON)
			}

			if err := w.Close(); err != nil {
				return
			}
		case <-ticker.C:
			c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
			if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
				return
			}
		}
	}
}

func main() {
	chatServer := NewChatServer()
	go chatServer.Start()

	// 静态文件处理
	fs := http.FileServer(http.Dir("./public"))
	http.Handle("/", fs)

	// WebSocket处理
	http.HandleFunc("/ws", chatServer.handleWebSocket)

	log.Println("服务器启动，监听端口 :8080")
	if err := http.ListenAndServe(":8080", nil); err != nil {
		log.Fatal("服务器启动失败:", err)
	}
}

// package main

// import (
// 	"encoding/json"
// 	"fmt"
// 	"log"
// 	"net/http"
// 	"time"

// 	"github.com/gorilla/websocket"
// )

// type Message struct {
// 	Username  string `json:"username"`
// 	Content   string `json:"content"`
// 	Type      string `json:"type"`
// 	Timestamp string `json:"timestamp"`
// }

// type Client struct {
// 	conn     *websocket.Conn
// 	username string
// 	send     chan Message
// }

// type ChatServer struct {
// 	clients    map[*Client]bool
// 	register   chan *Client
// 	unregister chan *Client
// 	broadcast  chan Message
// }

// // 创建带缓冲通道的聊天服务器
// func NewChatServer() *ChatServer {
// 	return &ChatServer{
// 		clients:    make(map[*Client]bool),
// 		register:   make(chan *Client), // 增加缓冲
// 		unregister: make(chan *Client), // 增加缓冲
// 		broadcast:  make(chan Message), // 增加大缓冲
// 	}
// }

// func (cs *ChatServer) Start() {
// 	for {
// 		select {
// 		case client := <-cs.register:
// 			cs.clients[client] = true
// 			log.Printf("客户端 %s 已连接\n", client.username)
// 			cs.broadcast <- Message{
// 				Username:  "系统",
// 				Content:   fmt.Sprintf("%s 加入了聊天", client.username),
// 				Type:      "system",
// 				Timestamp: time.Now().Format("15:04:05"),
// 			}
// 		case client := <-cs.unregister:
// 			if _, ok := cs.clients[client]; ok {
// 				delete(cs.clients, client)
// 				close(client.send)
// 				log.Printf("客户端 %s 已断开连接\n", client.username)
// 				cs.broadcast <- Message{
// 					Username:  "系统",
// 					Content:   fmt.Sprintf("%s 离开了聊天", client.username),
// 					Type:      "system",
// 					Timestamp: time.Now().Format("15:04:05"),
// 				}
// 			}
// 		case message := <-cs.broadcast:
// 			for client := range cs.clients {
// 				select {
// 				case client.send <- message: // 尝试非阻塞发送
// 				default:
// 					// 发送失败时关闭连接并移除
// 					close(client.send)
// 					delete(cs.clients, client)
// 					log.Printf("客户端 %s 因通道阻塞被移除\n", client.username)
// 				}
// 			}
// 		}
// 	}
// }

// var upgrader = websocket.Upgrader{
// 	ReadBufferSize:  1024,
// 	WriteBufferSize: 1024,
// 	CheckOrigin: func(r *http.Request) bool {
// 		return true
// 	},
// }

// func (cs *ChatServer) handleWebSocket(w http.ResponseWriter, r *http.Request) {
// 	username := r.URL.Query().Get("username")
// 	if username == "" {
// 		http.Error(w, "缺少用户名参数", http.StatusBadRequest)
// 		return
// 	}

// 	conn, err := upgrader.Upgrade(w, r, nil)
// 	if err != nil {
// 		log.Println("WebSocket升级失败:", err)
// 		return
// 	}

// 	client := &Client{
// 		conn:     conn,
// 		username: username,
// 		send:     make(chan Message, 256), // 带缓冲的发送通道
// 	}

// 	// 非阻塞注册
// 	select {
// 	case cs.register <- client:
// 		// 注册成功
// 	default:
// 		log.Println("注册通道已满，拒绝新连接")
// 		conn.Close()
// 		return
// 	}

// 	go client.writePump()
// 	go client.readPump(cs)
// }

// func (c *Client) readPump(cs *ChatServer) {
// 	defer func() {
// 		cs.unregister <- c
// 		c.conn.Close()
// 	}()

// 	c.conn.SetReadLimit(2048) // 增加消息大小限制
// 	c.conn.SetReadDeadline(time.Now().Add(60 * time.Second))
// 	c.conn.SetPongHandler(func(string) error {
// 		c.conn.SetReadDeadline(time.Now().Add(60 * time.Second))
// 		return nil
// 	})

// 	for {
// 		_, message, err := c.conn.ReadMessage()
// 		if err != nil {
// 			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
// 				log.Printf("读取错误: %v", err)
// 			}
// 			break
// 		}

// 		var msg Message
// 		if err := json.Unmarshal(message, &msg); err != nil {
// 			log.Printf("解析消息失败: %v", err)
// 			continue
// 		}

// 		msg.Username = c.username
// 		msg.Timestamp = time.Now().Format("15:04:05") // 添加时间戳

// 		cs.broadcast <- msg
// 	}
// }

// func (c *Client) writePump() {
// 	ticker := time.NewTicker(55 * time.Second)
// 	defer func() {
// 		ticker.Stop()
// 		c.conn.Close()
// 	}()

// 	for {
// 		select {
// 		case message, ok := <-c.send:
// 			c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
// 			if !ok {
// 				// 通道关闭
// 				c.conn.WriteMessage(websocket.CloseMessage, []byte{})
// 				return
// 			}

// 			messageJSON, err := json.Marshal(message)
// 			if err != nil {
// 				log.Printf("序列化消息失败: %v", err)
// 				continue
// 			}

// 			// 直接使用WriteMessage避免使用NextWriter
// 			if err := c.conn.WriteMessage(websocket.TextMessage, messageJSON); err != nil {
// 				log.Printf("写入失败: %v", err)
// 				return
// 			}
// 		case <-ticker.C:
// 			c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
// 			if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
// 				return
// 			}
// 		}
// 	}
// }

// func main() {
// 	chatServer := NewChatServer()
// 	go chatServer.Start()

// 	fs := http.FileServer(http.Dir("./public"))
// 	http.Handle("/", fs)
// 	http.HandleFunc("/ws", chatServer.handleWebSocket)

// 	log.Println("服务器启动，监听端口 :8080")
// 	if err := http.ListenAndServe(":8080", nil); err != nil {
// 		log.Fatal("服务器启动失败:", err)
// 	}
// }

// package main

// import (
// 	"encoding/json"
// 	"fmt"
// 	"log"
// 	"net/http"
// 	"sync"
// 	"sync/atomic"
// 	"time"

// 	"github.com/gorilla/websocket"
// )

// const (
// 	shardCount        = 32   // 分区数量，按CPU核心数调整
// 	maxConcurrentUser = 100000 // 最大并发用户数
// )

// type Message struct {
// 	Username  string `json:"username"`
// 	Content   string `json:"content"`
// 	Type      string `json:"type"`
// 	Timestamp string `json:"timestamp"`
// }

// type Client struct {
// 	conn     *websocket.Conn
// 	username string
// 	shardID  int
// 	lastSeen int64 // 原子操作的时间戳
// }

// type ChatShard struct {
// 	clients     map[*Client]bool
// 	clientMutex sync.RWMutex
// 	broadcast   chan Message
// }

// type ChatServer struct {
// 	shards   [shardCount]*ChatShard
// 	register chan *Client
// }

// func NewChatServer() *ChatServer {
// 	cs := &ChatServer{
// 		register: make(chan *Client, 1024),
// 	}

// 	// 初始化所有分区
// 	for i := 0; i < shardCount; i++ {
// 		cs.shards[i] = &ChatShard{
// 			clients:   make(map[*Client]bool),
// 			broadcast: make(chan Message, 1024),
// 		}
// 		go cs.shards[i].runShard()
// 	}

// 	return cs
// }

// func (shard *ChatShard) runShard() {
// 	for {
// 		select {
// 		case msg := <-shard.broadcast:
// 			shard.clientMutex.RLock()
// 			for client := range shard.clients {
// 				select {
// 				case client.send(msg):
// 				default:
// 					// 发送失败，关闭连接
// 					client.close()
// 				}
// 			}
// 			shard.clientMutex.RUnlock()
// 		}
// 	}
// }

// func (cs *ChatServer) Start() {
// 	// 注册处理
// 	for client := range cs.register {
// 		shardID := client.shardID
// 		shard := cs.shards[shardID]

// 		shard.clientMutex.Lock()
// 		if len(shard.clients) < maxConcurrentUser/shardCount {
// 			shard.clients[client] = true
// 			atomic.StoreInt64(&client.lastSeen, time.Now().UnixNano())
// 		}
// 		shard.clientMutex.Unlock()

// 		// 广播用户加入消息
// 		cs.Broadcast(Message{
// 			Username:  "系统",
// 			Content:   fmt.Sprintf("%s 加入了聊天", client.username),
// 			Type:      "system",
// 			Timestamp: time.Now().Format("15:04:05"),
// 		})
// 	}
// }

// // 向所有分区广播消息
// func (cs *ChatServer) Broadcast(msg Message) {
// 	for i := 0; i < shardCount; i++ {
// 		select {
// 		case cs.shards[i].broadcast <- msg:
// 		default:
// 			// 广播通道满，丢弃消息（或记录日志）
// 			log.Printf("警告: shard %d 广播队列已满", i)
// 		}
// 	}
// }

// var upgrader = websocket.Upgrader{
// 	ReadBufferSize:  1024,
// 	WriteBufferSize: 1024,
// 	CheckOrigin: func(r *http.Request) bool {
// 		return true
// 	},
// }

// func (cs *ChatServer) handleWebSocket(w http.ResponseWriter, r *http.Request) {
// 	username := r.URL.Query().Get("username")
// 	if username == "" {
// 		http.Error(w, "缺少用户名参数", http.StatusBadRequest)
// 		return
// 	}

// 	conn, err := upgrader.Upgrade(w, r, nil)
// 	if err != nil {
// 		log.Println("WebSocket升级失败:", err)
// 		return
// 	}

// 	// 使用一致性哈希分配分区
// 	shardID := int(fnv32(username) % shardCount)

// 	client := &Client{
// 		conn:     conn,
// 		username: username,
// 		shardID:  shardID,
// 	}

// 	// 非阻塞注册
// 	select {
// 	case cs.register <- client:
// 		// 注册成功
// 	default:
// 		log.Println("注册队列已满，拒绝新连接")
// 		conn.Close()
// 		return
// 	}

// 	go client.readPump(cs)
// }

// // 简单的FNV哈希用于分区分配
// func fnv32(key string) uint32 {
// 	hash := uint32(2166136261)
// 	const prime32 = uint32(16777619)
// 	for i := 0; i < len(key); i++ {
// 		hash *= prime32
// 		hash ^= uint32(key[i])
// 	}
// 	return hash
// }

// func (c *Client) send(msg Message) bool {
// 	messageJSON, err := json.Marshal(msg)
// 	if err != nil {
// 		log.Printf("序列化消息失败: %v", err)
// 		return false
// 	}

// 	c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
// 	if err := c.conn.WriteMessage(websocket.TextMessage, messageJSON); err != nil {
// 		return false
// 	}

// 	atomic.StoreInt64(&c.lastSeen, time.Now().UnixNano())
// 	return true
// }

// func (c *Client) close() {
// 	c.conn.Close()
// }

// func (c *Client) readPump(cs *ChatServer) {
// 	defer func() {
// 		// 注销客户端
// 		shard := cs.shards[c.shardID]
// 		shard.clientMutex.Lock()
// 		delete(shard.clients, c)
// 		shard.clientMutex.Unlock()

// 		c.close()

// 		// 广播用户离开消息
// 		cs.Broadcast(Message{
// 			Username:  "系统",
// 			Content:   fmt.Sprintf("%s 离开了聊天", c.username),
// 			Type:      "system",
// 			Timestamp: time.Now().Format("15:04:05"),
// 		})
// 	}()

// 	for {
// 		_, message, err := c.conn.ReadMessage()
// 		if err != nil {
// 			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
// 				log.Printf("读取错误: %v", err)
// 			}
// 			break
// 		}

// 		var msg Message
// 		if err := json.Unmarshal(message, &msg); err != nil {
// 			log.Printf("解析消息失败: %v", err)
// 			continue
// 		}

// 		msg.Username = c.username
// 		msg.Timestamp = time.Now().Format("15:04:05")
// 		atomic.StoreInt64(&c.lastSeen, time.Now().UnixNano())

// 		// 广播到所有分区
// 		cs.Broadcast(msg)
// 	}
// }

// func main() {
// 	chatServer := NewChatServer()
// 	go chatServer.Start()

// 	// 启动连接清理器
// 	go func() {
// 		ticker := time.NewTicker(5 * time.Minute)
// 		defer ticker.Stop()

// 		for range ticker.C {
// 			chatServer.cleanupStaleConnections()
// 		}
// 	}()

// 	fs := http.FileServer(http.Dir("./public"))
// 	http.Handle("/", fs)
// 	http.HandleFunc("/ws", chatServer.handleWebSocket)

// 	log.Println("服务器启动，监听端口 :8080")
// 	if err := http.ListenAndServe(":8080", nil); err != nil {
// 		log.Fatal("服务器启动失败:", err)
// 	}
// }

// // 清理空闲连接
// func (cs *ChatServer) cleanupStaleConnections() {
// 	timeout := time.Now().Add(-10 * time.Minute).UnixNano()

// 	for i := 0; i < shardCount; i++ {
// 		shard := cs.shards[i]
// 		shard.clientMutex.Lock()

// 		for client := range shard.clients {
// 			lastSeen := atomic.LoadInt64(&client.lastSeen)
// 			if lastSeen < timeout {
// 				client.close()
// 				delete(shard.clients, client)
// 				log.Printf("清理空闲连接: %s", client.username)
// 			}
// 		}

// 		shard.clientMutex.Unlock()
// 	}
// }
